# -*- coding: utf-8 -*-
# @Time    : 2023/9/23 18:39
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: serrulae-base


image_size = (1024, 1024)
max_steps = 64000

config = {
    "type": "Serrulae",
    "base_info": {
        "step": max_steps,
        "dot": 20,
        "save_iters": 1000,
        "pretrained": None,
        "save_path": "output/model/UNetPanV1_Base",
        "log_dir": "output/log_dir/UNetPanV1_Base",
        "only_last": True,
        "best_model": True,
    },
    "train_dataset": {
        "type": "Image2ImageDataset",
        "batch_size": 8,
        "shuffle": True,
        "num_workers": 4,
        "dt_root": "data/image",
        "gt_root": "data/groundtruth",
        "transforms": [
            {
                "type": "LoadData",
                "keys": ["img_1", "img_2"],
                "func": "cv2"
            },
            {
                "type": "Flip",
                "keys": ["img_1", "img_2"]
            },
            {
                "type": "ToTensor",
                "keys": ["img_1", "img_2"],
            },
        ]
    },
    "val_dataset": {
        "type": "Image2ImageDataset",
        "batch_size": 8,
        "shuffle": True,
        "num_workers": 4,
        "dt_root": "data/image",
        "gt_root": "data/groundtruth",
        "transforms": [
            {
                "type": "LoadData",
                "keys": ["img_1", "img_2"],
                "func": "cv2"
            },
            {
                "type": "ToTensor",
                "keys": ["img_1", "img_2"],
            }
        ]
    },
    "test_dataset": {
        "type": "Image2ImageDataset",
        "mode": "predict",
        "batch_size": 1,
        "shuffle": True,
        "num_workers": 2,
        "dt_root": "data/image",
        "save_root": "output/predict/UNetPanV1_NoNorm_conv_small_noise",
        "transforms": [
            {
                "type": "LoadData",
                "keys": ["img"],
                "func": "cv2"
            },
            {
                "type": "ToTensor",
                "keys": ["img"],
            }
        ]
    },
    "optimizer": {
        "type": "adam",
        "lr_scheduler": {
            "type": "WarmupCosineLR",    # Warm up 学习率刚开始是由小变大，Cosine
            "learning_rate": 0.0005,
            "total_steps": max_steps,
            "warmup_steps": 1000,
            "warmup_start_lr": 1e-8,
            "end_lr": 1e-8
        },
        "decay": None
    },
    "network": {
        "type": "image2image",
        "network": {
            "type": "UNetPanV1",
            "down_sample": "conv",
            "channels": [12, 24, 48, 96, 192]
        }
    },
    "loss": {
        "loss_list": [
            {
                "type": "MixLoss",
                "loss_list": [
                    {
                        "type": "L1Loss"
                    },
                    {
                        "type": "PSNRLoss"
                    }
                ],
                "loss_coef": [1, 0.02]
            }
        ],
        "loss_coef": [1]
    },
    "metric": [
        {
            "type": "PSNR",
            "coef": 0.005,
        },
        {
            "type": "SSIM",
            "coef": 0,
        },
        {
            "type": "MSSSIM",
            "coef": 0.5,
        }
    ]
}

